from ultralytics import YOLO
import cv2
import time
from ...Game.win import My_Window


# 加载训练好的模型
model = YOLO(model='D:/Project/All_Project/Python/blct/runs/train/cs8/weights/best.pt')
# 初始化截图模块
window_title = "雷电模拟器"  # 替换为你的窗口标题
capture =  My_Window(window_title)
# 实时推理函数
def real_time_inference():

    start_time = time.time()  # 开始时间
    frame_count = 0  # 帧计数器
    

    while True:
        # 获取截图
        frame = capture.capture()
        # 使用模型进行推理
        results = model(frame)
        # 解析推理结果
        for result in results:
            boxes = result.boxes.xyxy.cpu().numpy()  # 检测框坐标 (x1, y1, x2, y2)
            confidences = result.boxes.conf.cpu().numpy()  # 置信度
            class_ids = result.boxes.cls.cpu().numpy()  # 类别ID

            for box, confidence, class_id in zip(boxes, confidences, class_ids):
                x1, y1, x2, y2 = map(int, box)  # 转换为整数
                label = f"{model.names[int(class_id)]} {confidence:.2f}"  # 标签和置信度

                # 绘制检测框
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                # 绘制标签
                cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
                # 更新帧计数器和时间
                
        frame_count += 1
        current_time = time.time()
        elapsed_time = current_time - start_time

        # 每秒更新一次帧率
        if elapsed_time >= 1.0:
            fps = frame_count / elapsed_time
            print(f"FPS: {fps:.2f}")
            frame_count = 0
            start_time = current_time
        cv2.putText(frame, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        # 显示结果
        cv2.imshow("YOLO Real-Time Inference", frame)

        # 按 'q' 退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break



    # 释放资源
    cv2.destroyAllWindows()

# 启动实时推理
if __name__ == '__main__':
    real_time_inference()